import cv2
import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial.distance import cdist


def pick_symmetric_spots(image_path, roi=None,
                         inner_radius=15, outer_frac=0.4,
                         thres_rel=0.2, pixel_tol=2):
    """
    返回所有中心对称亮点的 (y, x) 坐标（相对于 ROI 中心）
    """
    # 1. 读图 & ROI
    img = cv2.imread(image_path, 0)
    if img is None:
        raise FileNotFoundError(image_path)
    if roi is None:
        roi = cv2.selectROI('select ROI', img)
        cv2.destroyAllWindows()
    x, y, w, h = roi
    roi_img = img[y:y + h, x:x + w]

    # 2. FFT + 中心化
    f = np.fft.fft2(roi_img)
    fshift = np.fft.fftshift(f)
    mag = np.abs(fshift)

    # 3. 环形掩码：去掉低频背景和高频噪声
    cy, cx = h // 2, w // 2
    Y, X = np.ogrid[:h, :w]
    dist = np.sqrt((X - cx) ** 2 + (Y - cy) ** 2)
    outer_radius = min(cx, cy) * outer_frac
    mask = (dist > inner_radius) & (dist < outer_radius)
    mag_band = mag * mask

    # 4. 二值化
    mag_norm = (mag_band - mag_band.min()) / (mag_band.max() - mag_band.min())
    _, bw = cv2.threshold(mag_norm.astype(np.float32), thres_rel, 1, cv2.THRESH_BINARY)

    # 5. 连通域 → 亮点坐标
    bw_u8 = (bw * 255).astype(np.uint8)
    n, labels, stats, centroids = cv2.connectedComponentsWithStats(bw_u8)
    # 去掉背景（stats 中面积最大的通常是背景）
    areas = stats[:, cv2.CC_STAT_AREA]
    idx = np.argsort(areas)[::-1][1:]  # 去掉最大面积
    pts = centroids[idx]  # [(x, y), ...]

    # 6. 对称配对
    center = np.array([cx, cy])
    pts_centered = pts - center
    pairs = []
    used = set()
    for i, p in enumerate(pts_centered):
        if i in used:
            continue
        # 对称点
        p_sym = -p
        dists = cdist([p_sym], pts_centered)[0]
        j = np.argmin(dists)
        if dists[j] < pixel_tol and j not in used:
            pairs.append((pts[i], pts[j]))
            used.update([i, j])

    # 7. 可视化
    plt.figure(figsize=(8, 4))
    plt.subplot(1, 2, 1)
    plt.title('Log magnitude')
    plt.imshow(20 * np.log(mag + 1), cmap='gray')
    plt.axis('off')

    plt.subplot(1, 2, 2)
    plt.title('Symmetric spots')
    plt.imshow(roi_img, cmap='gray')
    for (x1, y1), (x2, y2) in pairs:
        plt.plot([x1, x2], [y1, y2], 'r+', markersize=10)
    plt.axis('off')
    plt.tight_layout()
    plt.show()

    return pairs  # 每对两个点，顺序 (y, x)


# ---------------- demo ----------------
if __name__ == '__main__':
    pairs = pick_symmetric_spots('../img/pic3.jpg')
    print('共找到 %d 对对称亮点' % len(pairs))
    for (y1, x1), (y2, x2) in pairs:
        print(f'({y1:.1f},{x1:.1f}) ↔ ({y2:.1f},{x2:.1f})')
